{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# Merlin ETL, training and inference demo on the e-Commerce behavior data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "In this tutorial, we will be using the [eCommerce behavior data from multi category store](https://www.kaggle.com/mkechinov/ecommerce-behavior-data-from-multi-category-store) from [REES46 Marketing Platform](https://rees46.com/) as our dataset. This tutorial is built upon the NVIDIA RecSys 2020 [tutorial](https://recsys.acm.org/recsys20/tutorials/). \n",
    "\n",
    "This jupyter notebook provides the code to preprocess the dataset and generate the train, validation and test sets for the remainder of the tutorial. We define our own goal and filter the dataset accordingly.\n",
    "\n",
    "For our tutorial, we decided that our goal is to predict if a user purchased an item:\n",
    "\n",
    "-  Positive: User purchased an item\n",
    "-  Negative: User added an item to the cart, but did not purchase it (in the same session)    \n",
    "\n",
    "\n",
    "We split the dataset into train, validation and test set by the timestamp:\n",
    "- Training: October 2019 - February 2020\n",
    "- Validation: March 2020\n",
    "-  Test: April 2020\n",
    "\n",
    "We remove AddToCart Events from a session, if in the same session the same item was purchased.\n",
    "\n",
    "## Table of Contents\n",
    "1. [Data](#1)\n",
    "1. [ETL with NVTabular](#2)\n",
    "1. [Training with HugeCTR](#3)\n",
    "1. [HugeCTR inference](#4)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"1\"></a>\n",
    "## 1. Data\n",
    "First, we download and unzip the raw data.\n",
    "\n",
    "Note: the dataset is ~11GB and will take a while to download."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "export HOME=$PWD\n",
    "pip install gdown --user\n",
    "~/.local/bin/gdown  https://drive.google.com/uc?id=1-Rov9fFtGJqb7_ePc6qH-Rhzxn0cIcKB\n",
    "~/.local/bin/gdown  https://drive.google.com/uc?id=1-Rov9fFtGJqb7_ePc6qH-Rhzxn0cIcKB\n",
    "~/.local/bin/gdown  https://drive.google.com/uc?id=1zr_RXpGvOWN2PrWI6itWL8HnRsCpyqz8\n",
    "~/.local/bin/gdown  https://drive.google.com/uc?id=1g5WoIgLe05UMdREbxAjh0bEFgVCjA1UL\n",
    "~/.local/bin/gdown  https://drive.google.com/uc?id=1qZIwMbMgMmgDC5EoMdJ8aI9lQPsWA3-P\n",
    "~/.local/bin/gdown  https://drive.google.com/uc?id=1x5ohrrZNhWQN4Q-zww0RmXOwctKHH9PT\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['2020-Mar.csv.gz',\n",
       " '2020-Apr.csv.gz',\n",
       " '2019-Dec.csv.gz',\n",
       " '2020-Jan.csv.gz',\n",
       " '2020-Feb.csv.gz']"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import glob  \n",
    "\n",
    "list_files = glob.glob('*.csv.gz')\n",
    "list_files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data extraction and initial preprocessing\n",
    "\n",
    "We extract a few relevant columns from the raw datasets and parse date columns into several atomic colunns (day, month...)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "def process_files(file):\n",
    "    df_tmp = pd.read_csv(file, compression='gzip')\n",
    "    df_tmp['session_purchase'] =  df_tmp['user_session'] + '_' + df_tmp['product_id'].astype(str)\n",
    "    df_purchase = df_tmp[df_tmp['event_type']=='purchase']\n",
    "    df_cart = df_tmp[df_tmp['event_type']=='cart']\n",
    "    df_purchase = df_purchase[df_purchase['session_purchase'].isin(df_cart['session_purchase'])]\n",
    "    df_cart = df_cart[~(df_cart['session_purchase'].isin(df_purchase['session_purchase']))]\n",
    "    df_cart['target'] = 0\n",
    "    df_purchase['target'] = 1\n",
    "    df = pd.concat([df_cart, df_purchase])\n",
    "    df = df.drop('category_id', axis=1)\n",
    "    df = df.drop('session_purchase', axis=1)\n",
    "    df[['cat_0', 'cat_1', 'cat_2', 'cat_3']] = df['category_code'].str.split(\"\\.\", n = 3, expand = True).fillna('NA')\n",
    "    df['brand'] = df['brand'].fillna('NA')\n",
    "    df = df.drop('category_code', axis=1)\n",
    "    df['timestamp'] = pd.to_datetime(df['event_time'].str.replace(' UTC', ''))\n",
    "    df['ts_hour'] = df['timestamp'].dt.hour\n",
    "    df['ts_minute'] = df['timestamp'].dt.minute\n",
    "    df['ts_weekday'] = df['timestamp'].dt.weekday\n",
    "    df['ts_day'] = df['timestamp'].dt.day\n",
    "    df['ts_month'] = df['timestamp'].dt.month\n",
    "    df['ts_year'] = df['timestamp'].dt.year\n",
    "    df.to_csv('./dataset/' + file.replace('.gz', ''), index=False)\n",
    "    \n",
    "!mkdir ./dataset\n",
    "for file in tqdm(list_files):\n",
    "    print(file)\n",
    "    process_files(file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare train/validation/test data\n",
    "\n",
    "Next, we split the data into train, validation and test sets. We will be using 3 months for training, 1 month for validation and 1 month for testing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lp = []\n",
    "list_files = glob.glob('./dataset/*.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-rwxrwxrwx 1 26622 dip 479323170 Feb  1 01:04 ./dataset/2019-Dec.csv\n",
      "-rwxrwxrwx 1 26622 dip 455992639 Feb  1 01:04 ./dataset/2020-Apr.csv\n",
      "-rwxrwxrwx 1 26622 dip 453967664 Feb  1 01:05 ./dataset/2020-Feb.csv\n",
      "-rwxrwxrwx 1 26622 dip 375205173 Feb  1 01:04 ./dataset/2020-Jan.csv\n",
      "-rwxrwxrwx 1 26622 dip 403896607 Feb  1 01:05 ./dataset/2020-Mar.csv\n"
     ]
    }
   ],
   "source": [
    "!ls -l ./dataset/*.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "for file in list_files:\n",
    "    lp.append(pd.read_csv(file))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.concat(lp)\n",
    "df.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_test = df[df['ts_month']==4]\n",
    "df_valid = df[df['ts_month']==3]\n",
    "df_train = df[(df['ts_month']!=3)&(df['ts_month']!=4)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train.shape, df_valid.shape, df_test.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train.to_parquet('./data/train.parquet', index=False)\n",
    "df_valid.to_parquet('./data/valid.parquet', index=False)\n",
    "df_test.to_parquet('./data/test.parquet', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_train.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"2\"></a>\n",
    "## 2. Preprocessing with NVTabular\n",
    "\n",
    "Next, we will use NVTabular for preprocessing and engineering more features. \n",
    "\n",
    "But first, we need to import the necessary libraries and initialize a Dask GPU cluster for computation.\n",
    "\n",
    "### Initialize Dask GPU cluster\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.3.0+74.ge3ea5d0\n"
     ]
    }
   ],
   "source": [
    "# Standard Libraries\n",
    "import os\n",
    "from time import time\n",
    "import re\n",
    "import shutil\n",
    "import glob\n",
    "import warnings\n",
    "\n",
    "# External Dependencies\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import cupy as cp\n",
    "import cudf\n",
    "import dask_cudf\n",
    "from dask_cuda import LocalCUDACluster\n",
    "from dask.distributed import Client\n",
    "from dask.utils import parse_bytes\n",
    "from dask.delayed import delayed\n",
    "import rmm\n",
    "\n",
    "# NVTabular\n",
    "import nvtabular as nvt\n",
    "import nvtabular.ops as ops\n",
    "from nvtabular.io import Shuffle\n",
    "from nvtabular.utils import _pynvml_mem_size, device_mem_size\n",
    "\n",
    "print(nvt.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define some information about where to get our data\n",
    "BASE_DIR = \"./nvtabular_temp\"\n",
    "!rm -r $BASE_DIR && mkdir $BASE_DIR\n",
    "input_path = './dataset'\n",
    "dask_workdir = os.path.join(BASE_DIR, \"workdir\")\n",
    "output_path = os.path.join(BASE_DIR, \"output\")\n",
    "stats_path = os.path.join(BASE_DIR, \"stats\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example was tested on a DGX server with 8 GPUs. If you have less GPUs, modify the `NUM_GPUS` variable accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/rapids/lib/python3.8/site-packages/distributed/node.py:151: UserWarning: Port 8787 is already in use.\n",
      "Perhaps you already have a cluster running?\n",
      "Hosting the HTTP server on port 33387 instead\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table style=\"border: 2px solid white;\">\n",
       "<tr>\n",
       "<td style=\"vertical-align: top; border: 0px solid white\">\n",
       "<h3 style=\"text-align: left;\">Client</h3>\n",
       "<ul style=\"text-align: left; list-style: none; margin: 0; padding: 0;\">\n",
       "  <li><b>Scheduler: </b>tcp://127.0.0.1:41551</li>\n",
       "  <li><b>Dashboard: </b><a href='http://127.0.0.1:33387/status' target='_blank'>http://127.0.0.1:33387/status</a></li>\n",
       "</ul>\n",
       "</td>\n",
       "<td style=\"vertical-align: top; border: 0px solid white\">\n",
       "<h3 style=\"text-align: left;\">Cluster</h3>\n",
       "<ul style=\"text-align: left; list-style:none; margin: 0; padding: 0;\">\n",
       "  <li><b>Workers: </b>8</li>\n",
       "  <li><b>Cores: </b>8</li>\n",
       "  <li><b>Memory: </b>540.94 GB</li>\n",
       "</ul>\n",
       "</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "<Client: 'tcp://127.0.0.1:41551' processes=8 threads=8, memory=540.94 GB>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "NUM_GPUS = [0,1,2,3,4,5,6,7]\n",
    "#NUM_GPUS = [0]\n",
    "\n",
    "# Dask dashboard\n",
    "dashboard_port = \"8787\"\n",
    "\n",
    "# Deploy a Single-Machine Multi-GPU Cluster\n",
    "protocol = \"tcp\"             # \"tcp\" or \"ucx\"\n",
    "visible_devices = \",\".join([str(n) for n in NUM_GPUS])  # Delect devices to place workers\n",
    "device_limit_frac = 0.5      # Spill GPU-Worker memory to host at this limit.\n",
    "device_pool_frac = 0.6\n",
    "part_mem_frac = 0.05\n",
    "\n",
    "# Use total device size to calculate args.device_limit_frac\n",
    "device_size = device_mem_size(kind=\"total\")\n",
    "device_limit = int(device_limit_frac * device_size)\n",
    "device_pool_size = int(device_pool_frac * device_size)\n",
    "part_size = int(part_mem_frac * device_size)\n",
    "\n",
    "# Check if any device memory is already occupied\n",
    "\"\"\"\n",
    "for dev in visible_devices.split(\",\"):\n",
    "    fmem = _pynvml_mem_size(kind=\"free\", index=int(dev))\n",
    "    used = (device_size - fmem) / 1e9\n",
    "    if used > 1.0:\n",
    "        warnings.warn(f\"BEWARE - {used} GB is already occupied on device {int(dev)}!\")\n",
    "\"\"\"\n",
    "\n",
    "cluster = None               # (Optional) Specify existing scheduler port\n",
    "if cluster is None:\n",
    "    cluster = LocalCUDACluster(\n",
    "        protocol = protocol,\n",
    "        n_workers=len(visible_devices.split(\",\")),\n",
    "        CUDA_VISIBLE_DEVICES = visible_devices,\n",
    "        device_memory_limit = device_limit,\n",
    "        local_directory=dask_workdir,\n",
    "        dashboard_address=\":\" + dashboard_port,\n",
    "    )\n",
    "\n",
    "# Create the distributed client\n",
    "client = Client(cluster)\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tue Mar  2 10:56:23 2021       \n",
      "+-----------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 450.80.02    Driver Version: 450.80.02    CUDA Version: 11.0     |\n",
      "|-------------------------------+----------------------+----------------------+\n",
      "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                               |                      |               MIG M. |\n",
      "|===============================+======================+======================|\n",
      "|   0  Tesla V100-SXM2...  On   | 00000000:06:00.0 Off |                    0 |\n",
      "| N/A   34C    P0    57W / 300W |   5618MiB / 32510MiB |      0%      Default |\n",
      "|                               |                      |                  N/A |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|   1  Tesla V100-SXM2...  On   | 00000000:07:00.0 Off |                    0 |\n",
      "| N/A   35C    P0    56W / 300W |    613MiB / 32510MiB |      0%      Default |\n",
      "|                               |                      |                  N/A |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|   2  Tesla V100-SXM2...  On   | 00000000:0A:00.0 Off |                    0 |\n",
      "| N/A   36C    P0    56W / 300W |    613MiB / 32510MiB |      0%      Default |\n",
      "|                               |                      |                  N/A |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|   3  Tesla V100-SXM2...  On   | 00000000:0B:00.0 Off |                    0 |\n",
      "| N/A   32C    P0    56W / 300W |    613MiB / 32510MiB |      0%      Default |\n",
      "|                               |                      |                  N/A |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|   4  Tesla V100-SXM2...  On   | 00000000:85:00.0 Off |                    0 |\n",
      "| N/A   34C    P0    56W / 300W |    613MiB / 32510MiB |      0%      Default |\n",
      "|                               |                      |                  N/A |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|   5  Tesla V100-SXM2...  On   | 00000000:86:00.0 Off |                    0 |\n",
      "| N/A   34C    P0    56W / 300W |    613MiB / 32510MiB |      0%      Default |\n",
      "|                               |                      |                  N/A |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|   6  Tesla V100-SXM2...  On   | 00000000:89:00.0 Off |                    0 |\n",
      "| N/A   36C    P0    56W / 300W |    613MiB / 32510MiB |      0%      Default |\n",
      "|                               |                      |                  N/A |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "|   7  Tesla V100-SXM2...  On   | 00000000:8A:00.0 Off |                    0 |\n",
      "| N/A   38C    P0    55W / 300W |    613MiB / 32510MiB |      0%      Default |\n",
      "|                               |                      |                  N/A |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "                                                                               \n",
      "+-----------------------------------------------------------------------------+\n",
      "| Processes:                                                                  |\n",
      "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
      "|        ID   ID                                                   Usage      |\n",
      "|=============================================================================|\n",
      "+-----------------------------------------------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'tcp://127.0.0.1:34701': None,\n",
       " 'tcp://127.0.0.1:35083': None,\n",
       " 'tcp://127.0.0.1:37071': None,\n",
       " 'tcp://127.0.0.1:38719': None,\n",
       " 'tcp://127.0.0.1:42111': None,\n",
       " 'tcp://127.0.0.1:43367': None,\n",
       " 'tcp://127.0.0.1:46323': None,\n",
       " 'tcp://127.0.0.1:46865': None}"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Initialize RMM pool on ALL workers\n",
    "def _rmm_pool():\n",
    "    rmm.reinitialize(\n",
    "        # RMM may require the pool size to be a multiple of 256.\n",
    "        pool_allocator=True,\n",
    "        initial_pool_size=(device_pool_size // 256) * 256, # Use default size\n",
    "    )\n",
    "    \n",
    "client.run(_rmm_pool)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define NVTabular dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_paths = glob.glob('./data/train.parquet')\n",
    "valid_paths = glob.glob('./data/valid.parquet')\n",
    "test_paths = glob.glob('./data/test.parquet')\n",
    "\n",
    "train_dataset = nvt.Dataset(train_paths, engine='parquet', part_mem_fraction=0.15)\n",
    "valid_dataset = nvt.Dataset(valid_paths, engine='parquet', part_mem_fraction=0.15)\n",
    "test_dataset = nvt.Dataset(test_paths, engine='parquet', part_mem_fraction=0.15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>event_time</th>\n",
       "      <th>event_type</th>\n",
       "      <th>product_id</th>\n",
       "      <th>brand</th>\n",
       "      <th>price</th>\n",
       "      <th>user_id</th>\n",
       "      <th>user_session</th>\n",
       "      <th>target</th>\n",
       "      <th>cat_0</th>\n",
       "      <th>cat_1</th>\n",
       "      <th>cat_2</th>\n",
       "      <th>cat_3</th>\n",
       "      <th>timestamp</th>\n",
       "      <th>ts_hour</th>\n",
       "      <th>ts_minute</th>\n",
       "      <th>ts_weekday</th>\n",
       "      <th>ts_day</th>\n",
       "      <th>ts_month</th>\n",
       "      <th>ts_year</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2020-02-01 00:00:18 UTC</td>\n",
       "      <td>cart</td>\n",
       "      <td>100065078</td>\n",
       "      <td>xiaomi</td>\n",
       "      <td>568.61</td>\n",
       "      <td>526615078</td>\n",
       "      <td>5f0aab9f-f92e-4eff-b0d2-fcec5f553f01</td>\n",
       "      <td>0</td>\n",
       "      <td>construction</td>\n",
       "      <td>tools</td>\n",
       "      <td>light</td>\n",
       "      <td>&lt;NA&gt;</td>\n",
       "      <td>2020-02-01 00:00:18</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>5</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>2020</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2020-02-01 00:00:18 UTC</td>\n",
       "      <td>cart</td>\n",
       "      <td>5701246</td>\n",
       "      <td>&lt;NA&gt;</td>\n",
       "      <td>24.43</td>\n",
       "      <td>563902689</td>\n",
       "      <td>76cc9152-8a9f-43e9-b98a-ee484510f379</td>\n",
       "      <td>0</td>\n",
       "      <td>electronics</td>\n",
       "      <td>video</td>\n",
       "      <td>tv</td>\n",
       "      <td>&lt;NA&gt;</td>\n",
       "      <td>2020-02-01 00:00:18</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>5</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>2020</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2020-02-01 00:00:31 UTC</td>\n",
       "      <td>cart</td>\n",
       "      <td>14701533</td>\n",
       "      <td>&lt;NA&gt;</td>\n",
       "      <td>154.42</td>\n",
       "      <td>520953435</td>\n",
       "      <td>5f1c7752-cf92-41fc-9a16-e8897a90eee8</td>\n",
       "      <td>0</td>\n",
       "      <td>electronics</td>\n",
       "      <td>video</td>\n",
       "      <td>projector</td>\n",
       "      <td>&lt;NA&gt;</td>\n",
       "      <td>2020-02-01 00:00:31</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>5</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>2020</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2020-02-01 00:00:40 UTC</td>\n",
       "      <td>cart</td>\n",
       "      <td>1004855</td>\n",
       "      <td>xiaomi</td>\n",
       "      <td>123.30</td>\n",
       "      <td>519236281</td>\n",
       "      <td>e512f514-dc7f-4fc9-9042-e3955989d395</td>\n",
       "      <td>0</td>\n",
       "      <td>construction</td>\n",
       "      <td>tools</td>\n",
       "      <td>light</td>\n",
       "      <td>&lt;NA&gt;</td>\n",
       "      <td>2020-02-01 00:00:40</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>5</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>2020</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2020-02-01 00:00:47 UTC</td>\n",
       "      <td>cart</td>\n",
       "      <td>1005100</td>\n",
       "      <td>samsung</td>\n",
       "      <td>140.28</td>\n",
       "      <td>550305600</td>\n",
       "      <td>bd7a37b6-420d-4575-8852-ac825aff39b5</td>\n",
       "      <td>0</td>\n",
       "      <td>construction</td>\n",
       "      <td>tools</td>\n",
       "      <td>light</td>\n",
       "      <td>&lt;NA&gt;</td>\n",
       "      <td>2020-02-01 00:00:47</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>5</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>2020</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                event_time event_type  product_id    brand   price    user_id  \\\n",
       "0  2020-02-01 00:00:18 UTC       cart   100065078   xiaomi  568.61  526615078   \n",
       "1  2020-02-01 00:00:18 UTC       cart     5701246     <NA>   24.43  563902689   \n",
       "2  2020-02-01 00:00:31 UTC       cart    14701533     <NA>  154.42  520953435   \n",
       "3  2020-02-01 00:00:40 UTC       cart     1004855   xiaomi  123.30  519236281   \n",
       "4  2020-02-01 00:00:47 UTC       cart     1005100  samsung  140.28  550305600   \n",
       "\n",
       "                           user_session  target         cat_0  cat_1  \\\n",
       "0  5f0aab9f-f92e-4eff-b0d2-fcec5f553f01       0  construction  tools   \n",
       "1  76cc9152-8a9f-43e9-b98a-ee484510f379       0   electronics  video   \n",
       "2  5f1c7752-cf92-41fc-9a16-e8897a90eee8       0   electronics  video   \n",
       "3  e512f514-dc7f-4fc9-9042-e3955989d395       0  construction  tools   \n",
       "4  bd7a37b6-420d-4575-8852-ac825aff39b5       0  construction  tools   \n",
       "\n",
       "       cat_2 cat_3            timestamp  ts_hour  ts_minute  ts_weekday  \\\n",
       "0      light  <NA>  2020-02-01 00:00:18        0          0           5   \n",
       "1         tv  <NA>  2020-02-01 00:00:18        0          0           5   \n",
       "2  projector  <NA>  2020-02-01 00:00:31        0          0           5   \n",
       "3      light  <NA>  2020-02-01 00:00:40        0          0           5   \n",
       "4      light  <NA>  2020-02-01 00:00:47        0          0           5   \n",
       "\n",
       "   ts_day  ts_month  ts_year  \n",
       "0       1         2     2020  \n",
       "1       1         2     2020  \n",
       "2       1         2     2020  \n",
       "3       1         2     2020  \n",
       "4       1         2     2020  "
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset.to_ddf().head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "19"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train_dataset.to_ddf().columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Index(['event_time', 'event_type', 'product_id', 'brand', 'price', 'user_id',\n",
       "       'user_session', 'target', 'cat_0', 'cat_1', 'cat_2', 'cat_3',\n",
       "       'timestamp', 'ts_hour', 'ts_minute', 'ts_weekday', 'ts_day', 'ts_month',\n",
       "       'ts_year'],\n",
       "      dtype='object')"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset.to_ddf().columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "7949839"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train_dataset.to_ddf())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preprocessing and feature engineering"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook we will explore a few feature engineering technique with NVTabular:\n",
    "\n",
    "- Creating cross features, e.g. `user_id` and `'brand`\n",
    "- Target encoding\n",
    "\n",
    "The engineered features will then be preprocessed into a form suitable for machine learning model:\n",
    "\n",
    "- Fill missing values\n",
    "- Encoding categorical features into integer values\n",
    "- Normalization of numeric features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nvtabular.ops import LambdaOp\n",
    "\n",
    "# cross features\n",
    "def user_id_cross_maker(col, gdf):\n",
    "    return col.astype(str) + '_' + gdf['user_id'].astype(str)\n",
    "\n",
    "user_id_cross_features = (\n",
    "    nvt.ColumnGroup(['product_id', 'brand', 'ts_hour', 'ts_minute']) >>\n",
    "    LambdaOp(user_id_cross_maker, dependency=['user_id']) >> \n",
    "    nvt.ops.Rename(postfix = '_user_id_cross')\n",
    ")\n",
    "\n",
    "\n",
    "def user_id_brand_cross_maker(col, gdf):\n",
    "    return col.astype(str) + '_' + gdf['user_id'].astype(str) + '_' + gdf['brand'].astype(str)\n",
    "\n",
    "user_id_brand_cross_features = (\n",
    "    nvt.ColumnGroup(['ts_hour', 'ts_weekday', 'cat_0', 'cat_1', 'cat_2']) >>\n",
    "    LambdaOp(user_id_brand_cross_maker, dependency=['user_id', 'brand']) >> \n",
    "    nvt.ops.Rename(postfix = '_user_id_brand_cross')\n",
    ")\n",
    "\n",
    "target_encode = (\n",
    "    ['brand', 'user_id', 'product_id', 'cat_2', ['ts_weekday', 'ts_day']] >>\n",
    "    nvt.ops.TargetEncoding(\n",
    "        nvt.ColumnGroup('target'),\n",
    "        kfold=5,\n",
    "        p_smooth=20,\n",
    "        out_dtype=\"float32\",\n",
    "        )\n",
    ")\n",
    "\n",
    "cat_feats = (user_id_brand_cross_features + user_id_cross_features) >> nvt.ops.Categorify()\n",
    "cont_feats =  ['price', 'ts_weekday', 'ts_day', 'ts_month'] >> nvt.ops.FillMissing() >>  nvt.ops.Normalize()\n",
    "cont_feats += target_encode >> nvt.ops.Rename(postfix = '_TE')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = cat_feats + cont_feats + 'target'\n",
    "proc = nvt.Workflow(output)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize workflow as a DAG\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reading package lists... Done\n",
      "Building dependency tree       \n",
      "Reading state information... Done\n",
      "graphviz is already the newest version (2.42.2-3build2).\n",
      "0 upgraded, 0 newly installed, 0 to remove and 9 not upgraded.\n"
     ]
    }
   ],
   "source": [
    "!apt install graphviz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 2.43.0 (0)\n",
       " -->\n",
       "<!-- Title: %3 Pages: 1 -->\n",
       "<svg width=\"1554pt\" height=\"476pt\"\n",
       " viewBox=\"0.00 0.00 1553.96 476.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 472)\">\n",
       "<title>%3</title>\n",
       "<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-472 1549.96,-472 1549.96,4 -4,4\"/>\n",
       "<!-- 0 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>0</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"484.48\" cy=\"-306\" rx=\"194.97\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"484.48\" y=\"-302.3\" font-family=\"Times,serif\" font-size=\"14.00\">input cols=[price, ts_weekday, ts_day...]</text>\n",
       "</g>\n",
       "<!-- 3 -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>3</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"546.48\" cy=\"-234\" rx=\"62.29\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"546.48\" y=\"-230.3\" font-family=\"Times,serif\" font-size=\"14.00\">FillMissing</text>\n",
       "</g>\n",
       "<!-- 0&#45;&gt;3 -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>0&#45;&gt;3</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M499.81,-287.7C507.43,-279.09 516.76,-268.55 525.08,-259.17\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"527.75,-261.43 531.76,-251.62 522.51,-256.79 527.75,-261.43\"/>\n",
       "</g>\n",
       "<!-- 1 -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>1</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"865.48\" cy=\"-90\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"865.48\" y=\"-86.3\" font-family=\"Times,serif\" font-size=\"14.00\">+</text>\n",
       "</g>\n",
       "<!-- 20 -->\n",
       "<g id=\"node21\" class=\"node\">\n",
       "<title>20</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"865.48\" cy=\"-18\" rx=\"505.81\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"865.48\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\">output cols=[ts_hour_user_id_brand_cross, ts_weekday_user_id_brand_cross, cat_0_user_id_brand_cross...]</text>\n",
       "</g>\n",
       "<!-- 1&#45;&gt;20 -->\n",
       "<g id=\"edge20\" class=\"edge\">\n",
       "<title>1&#45;&gt;20</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M865.48,-71.7C865.48,-63.98 865.48,-54.71 865.48,-46.11\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"868.98,-46.1 865.48,-36.1 861.98,-46.1 868.98,-46.1\"/>\n",
       "</g>\n",
       "<!-- 17 -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>17</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"802.48\" cy=\"-162\" rx=\"59.59\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"802.48\" y=\"-158.3\" font-family=\"Times,serif\" font-size=\"14.00\">Categorify</text>\n",
       "</g>\n",
       "<!-- 17&#45;&gt;1 -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>17&#45;&gt;1</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M817.41,-144.41C825.68,-135.22 836.08,-123.67 845.09,-113.66\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"847.86,-115.81 851.95,-106.04 842.66,-111.13 847.86,-115.81\"/>\n",
       "</g>\n",
       "<!-- 5 -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>5</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"635.48\" cy=\"-162\" rx=\"58.49\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"635.48\" y=\"-158.3\" font-family=\"Times,serif\" font-size=\"14.00\">Normalize</text>\n",
       "</g>\n",
       "<!-- 5&#45;&gt;1 -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>5&#45;&gt;1</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M675.55,-148.81C719.78,-135.35 790.2,-113.91 831.73,-101.27\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"832.84,-104.59 841.38,-98.33 830.8,-97.9 832.84,-104.59\"/>\n",
       "</g>\n",
       "<!-- 7 -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>7</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"928.48\" cy=\"-162\" rx=\"48.19\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"928.48\" y=\"-158.3\" font-family=\"Times,serif\" font-size=\"14.00\">Rename</text>\n",
       "</g>\n",
       "<!-- 7&#45;&gt;1 -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>7&#45;&gt;1</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M913.87,-144.76C905.41,-135.37 894.65,-123.41 885.42,-113.15\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"888.02,-110.81 878.73,-105.72 882.82,-115.49 888.02,-110.81\"/>\n",
       "</g>\n",
       "<!-- 15 -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>15</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"1094.48\" cy=\"-162\" rx=\"99.38\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"1094.48\" y=\"-158.3\" font-family=\"Times,serif\" font-size=\"14.00\">input cols=[target]</text>\n",
       "</g>\n",
       "<!-- 15&#45;&gt;1 -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>15&#45;&gt;1</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1046.2,-146.24C1002.11,-132.76 938.07,-113.19 899.24,-101.32\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"900.06,-97.91 889.47,-98.33 898.01,-104.6 900.06,-97.91\"/>\n",
       "</g>\n",
       "<!-- 2 -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>2</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"745.48\" cy=\"-306\" rx=\"48.19\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"745.48\" y=\"-302.3\" font-family=\"Times,serif\" font-size=\"14.00\">Rename</text>\n",
       "</g>\n",
       "<!-- 19 -->\n",
       "<g id=\"node20\" class=\"node\">\n",
       "<title>19</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"802.48\" cy=\"-234\" rx=\"27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"802.48\" y=\"-230.3\" font-family=\"Times,serif\" font-size=\"14.00\">+</text>\n",
       "</g>\n",
       "<!-- 2&#45;&gt;19 -->\n",
       "<g id=\"edge18\" class=\"edge\">\n",
       "<title>2&#45;&gt;19</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M758.99,-288.41C766.29,-279.44 775.44,-268.21 783.45,-258.37\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"786.25,-260.48 789.85,-250.51 780.82,-256.06 786.25,-260.48\"/>\n",
       "</g>\n",
       "<!-- 10 -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>10</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"636.48\" cy=\"-378\" rx=\"61.19\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"636.48\" y=\"-374.3\" font-family=\"Times,serif\" font-size=\"14.00\">LambdaOp</text>\n",
       "</g>\n",
       "<!-- 10&#45;&gt;2 -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>10&#45;&gt;2</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M660.67,-361.46C676.27,-351.45 696.72,-338.32 713.63,-327.46\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"715.58,-330.36 722.11,-322.01 711.8,-324.47 715.58,-330.36\"/>\n",
       "</g>\n",
       "<!-- 3&#45;&gt;5 -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>3&#45;&gt;5</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M567.12,-216.76C579,-207.42 594.1,-195.55 607.09,-185.33\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"609.56,-187.84 615.26,-178.91 605.23,-182.34 609.56,-187.84\"/>\n",
       "</g>\n",
       "<!-- 4 -->\n",
       "<g id=\"node10\" class=\"node\">\n",
       "<title>4</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"967.48\" cy=\"-450\" rx=\"104.78\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"967.48\" y=\"-446.3\" font-family=\"Times,serif\" font-size=\"14.00\">input cols=[user_id]</text>\n",
       "</g>\n",
       "<!-- 6 -->\n",
       "<g id=\"node11\" class=\"node\">\n",
       "<title>6</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"967.48\" cy=\"-378\" rx=\"61.19\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"967.48\" y=\"-374.3\" font-family=\"Times,serif\" font-size=\"14.00\">LambdaOp</text>\n",
       "</g>\n",
       "<!-- 4&#45;&gt;6 -->\n",
       "<g id=\"edge9\" class=\"edge\">\n",
       "<title>4&#45;&gt;6</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M967.48,-431.7C967.48,-423.98 967.48,-414.71 967.48,-406.11\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"970.98,-406.1 967.48,-396.1 963.98,-406.1 970.98,-406.1\"/>\n",
       "</g>\n",
       "<!-- 11 -->\n",
       "<g id=\"node18\" class=\"node\">\n",
       "<title>11</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"859.48\" cy=\"-306\" rx=\"48.19\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"859.48\" y=\"-302.3\" font-family=\"Times,serif\" font-size=\"14.00\">Rename</text>\n",
       "</g>\n",
       "<!-- 6&#45;&gt;11 -->\n",
       "<g id=\"edge14\" class=\"edge\">\n",
       "<title>6&#45;&gt;11</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M943.24,-361.29C927.92,-351.36 907.95,-338.41 891.35,-327.66\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"892.82,-324.44 882.53,-321.94 889.02,-330.31 892.82,-324.44\"/>\n",
       "</g>\n",
       "<!-- 18 -->\n",
       "<g id=\"node12\" class=\"node\">\n",
       "<title>18</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"1291.48\" cy=\"-450\" rx=\"201.46\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"1291.48\" y=\"-446.3\" font-family=\"Times,serif\" font-size=\"14.00\">input cols=[product_id, brand, ts_hour...]</text>\n",
       "</g>\n",
       "<!-- 18&#45;&gt;6 -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>18&#45;&gt;6</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1218.37,-433.2C1159.98,-420.59 1079.28,-403.15 1025.33,-391.5\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1025.76,-388.01 1015.25,-389.32 1024.28,-394.85 1025.76,-388.01\"/>\n",
       "</g>\n",
       "<!-- 12 -->\n",
       "<g id=\"node13\" class=\"node\">\n",
       "<title>12</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"1025.48\" cy=\"-234\" rx=\"84.49\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"1025.48\" y=\"-230.3\" font-family=\"Times,serif\" font-size=\"14.00\">TargetEncoding</text>\n",
       "</g>\n",
       "<!-- 12&#45;&gt;7 -->\n",
       "<g id=\"edge10\" class=\"edge\">\n",
       "<title>12&#45;&gt;7</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1002.49,-216.41C989.23,-206.84 972.4,-194.69 958.14,-184.4\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"959.8,-181.29 949.65,-178.27 955.71,-186.96 959.8,-181.29\"/>\n",
       "</g>\n",
       "<!-- 8 -->\n",
       "<g id=\"node14\" class=\"node\">\n",
       "<title>8</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"1025.48\" cy=\"-306\" rx=\"99.38\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"1025.48\" y=\"-302.3\" font-family=\"Times,serif\" font-size=\"14.00\">input cols=[target]</text>\n",
       "</g>\n",
       "<!-- 8&#45;&gt;12 -->\n",
       "<g id=\"edge16\" class=\"edge\">\n",
       "<title>8&#45;&gt;12</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1025.48,-287.7C1025.48,-279.98 1025.48,-270.71 1025.48,-262.11\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1028.98,-262.1 1025.48,-252.1 1021.98,-262.1 1028.98,-262.1\"/>\n",
       "</g>\n",
       "<!-- 9 -->\n",
       "<g id=\"node15\" class=\"node\">\n",
       "<title>9</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"201.48\" cy=\"-450\" rx=\"201.46\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"201.48\" y=\"-446.3\" font-family=\"Times,serif\" font-size=\"14.00\">input cols=[ts_hour, ts_weekday, cat_0...]</text>\n",
       "</g>\n",
       "<!-- 9&#45;&gt;10 -->\n",
       "<g id=\"edge11\" class=\"edge\">\n",
       "<title>9&#45;&gt;10</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M294.53,-434.03C378.92,-420.45 500.97,-400.81 574.14,-389.03\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"574.81,-392.47 584.13,-387.42 573.7,-385.56 574.81,-392.47\"/>\n",
       "</g>\n",
       "<!-- 14 -->\n",
       "<g id=\"node16\" class=\"node\">\n",
       "<title>14</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"525.48\" cy=\"-450\" rx=\"104.78\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"525.48\" y=\"-446.3\" font-family=\"Times,serif\" font-size=\"14.00\">input cols=[user_id]</text>\n",
       "</g>\n",
       "<!-- 14&#45;&gt;10 -->\n",
       "<g id=\"edge12\" class=\"edge\">\n",
       "<title>14&#45;&gt;10</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M551.79,-432.41C567.15,-422.72 586.68,-410.4 603.12,-400.04\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"605.29,-402.81 611.89,-394.51 601.56,-396.88 605.29,-402.81\"/>\n",
       "</g>\n",
       "<!-- 16 -->\n",
       "<g id=\"node17\" class=\"node\">\n",
       "<title>16</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"746.48\" cy=\"-450\" rx=\"98.28\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"746.48\" y=\"-446.3\" font-family=\"Times,serif\" font-size=\"14.00\">input cols=[brand]</text>\n",
       "</g>\n",
       "<!-- 16&#45;&gt;10 -->\n",
       "<g id=\"edge13\" class=\"edge\">\n",
       "<title>16&#45;&gt;10</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M720.69,-432.59C705.48,-422.9 686.06,-410.55 669.7,-400.14\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"671.29,-397 660.98,-394.59 667.54,-402.91 671.29,-397\"/>\n",
       "</g>\n",
       "<!-- 11&#45;&gt;19 -->\n",
       "<g id=\"edge19\" class=\"edge\">\n",
       "<title>11&#45;&gt;19</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M845.97,-288.41C838.67,-279.44 829.53,-268.21 821.51,-258.37\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"824.14,-256.06 815.11,-250.51 818.71,-260.48 824.14,-256.06\"/>\n",
       "</g>\n",
       "<!-- 13 -->\n",
       "<g id=\"node19\" class=\"node\">\n",
       "<title>13</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"1344.48\" cy=\"-306\" rx=\"201.46\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"1344.48\" y=\"-302.3\" font-family=\"Times,serif\" font-size=\"14.00\">input cols=[brand, user_id, product_id...]</text>\n",
       "</g>\n",
       "<!-- 13&#45;&gt;12 -->\n",
       "<g id=\"edge15\" class=\"edge\">\n",
       "<title>13&#45;&gt;12</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M1272.1,-289.12C1217.96,-277.24 1144.5,-261.12 1091.86,-249.57\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"1092.48,-246.12 1081.96,-247.39 1090.97,-252.95 1092.48,-246.12\"/>\n",
       "</g>\n",
       "<!-- 19&#45;&gt;17 -->\n",
       "<g id=\"edge17\" class=\"edge\">\n",
       "<title>19&#45;&gt;17</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M802.48,-215.7C802.48,-207.98 802.48,-198.71 802.48,-190.11\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"805.98,-190.1 802.48,-180.1 798.98,-190.1 805.98,-190.1\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.dot.Digraph at 0x7eff343490a0>"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output.graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Executing the workflow\n",
    "\n",
    "After having defined the workflow, calling the `fit()` method will start the actual computation to record the required statistics from the training data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 18.5 s, sys: 3.47 s, total: 22 s\n",
      "Wall time: 41.4 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "time_preproc_start = time()\n",
    "proc.fit(train_dataset)\n",
    "time_preproc = time()-time_preproc_start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "dict_dtypes = {}\n",
    "for col in cat_feats.columns:\n",
    "    dict_dtypes[col] = np.int64\n",
    "for col in cont_feats.columns:\n",
    "    dict_dtypes[col] = np.float32\n",
    "\n",
    "dict_dtypes['target'] = np.float32\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we call the `transform()` method to transform the datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_train_dir = os.path.join(output_path, 'train/')\n",
    "output_valid_dir = os.path.join(output_path, 'valid/')\n",
    "output_test_dir = os.path.join(output_path, 'test/')\n",
    "! rm -rf $output_train_dir && mkdir -p $output_train_dir\n",
    "! rm -rf $output_valid_dir && mkdir -p $output_valid_dir\n",
    "! rm -rf $output_test_dir && mkdir -p $output_test_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.23 s, sys: 1.9 s, total: 4.13 s\n",
      "Wall time: 10.2 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "time_preproc_start = time()\n",
    "proc.transform(train_dataset).to_parquet(output_path=output_train_dir, dtypes=dict_dtypes,\n",
    "                                         shuffle=nvt.io.Shuffle.PER_PARTITION,\n",
    "                                         cats=cat_feats.columns,\n",
    "                                         conts=cont_feats.columns,\n",
    "                                         labels=['target'])\n",
    "time_preproc += time()-time_preproc_start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 692580\n",
      "-rw-r--r-- 1 nobody nogroup 706376015 Mar  2 11:09 0.966e31e2c9cb4610b5adf01a1ce5f39f.parquet\n",
      "-rw-r--r-- 1 nobody nogroup        75 Mar  2 11:09 _file_list.txt\n",
      "-rw-r--r-- 1 nobody nogroup     21931 Mar  2 11:09 _metadata\n",
      "-rw-r--r-- 1 nobody nogroup      1073 Mar  2 11:09 _metadata.json\n"
     ]
    }
   ],
   "source": [
    "!ls -l $output_train_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 975 ms, sys: 703 ms, total: 1.68 s\n",
      "Wall time: 2.97 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "time_preproc_start = time()\n",
    "proc.transform(valid_dataset).to_parquet(output_path=output_valid_dir, dtypes=dict_dtypes,\n",
    "                                         shuffle=nvt.io.Shuffle.PER_PARTITION,\n",
    "                                         cats=cat_feats.columns,\n",
    "                                         conts=cont_feats.columns,\n",
    "                                         labels=['target'])\n",
    "time_preproc += time()-time_preproc_start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 90816\n",
      "-rw-r--r-- 1 nobody nogroup 92605522 Mar  2 11:09 0.a0a34fa76bd3494fbc29f1c0fdbb5b69.parquet\n",
      "-rw-r--r-- 1 nobody nogroup       75 Mar  2 11:09 _file_list.txt\n",
      "-rw-r--r-- 1 nobody nogroup    10351 Mar  2 11:09 _metadata\n",
      "-rw-r--r-- 1 nobody nogroup     1073 Mar  2 11:09 _metadata.json\n"
     ]
    }
   ],
   "source": [
    "!ls -l $output_valid_dir"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 816 ms, sys: 851 ms, total: 1.67 s\n",
      "Wall time: 2.2 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "time_preproc_start = time()\n",
    "proc.transform(test_dataset).to_parquet(output_path=output_test_dir, dtypes=dict_dtypes,\n",
    "                                         shuffle=nvt.io.Shuffle.PER_PARTITION,\n",
    "                                         cats=cat_feats.columns,\n",
    "                                         conts=cont_feats.columns,\n",
    "                                         labels=['target'])\n",
    "time_preproc += time()-time_preproc_start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "56.746562004089355"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "time_preproc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Verify the preprocessed data\n",
    "\n",
    "Let's quickly read the data back and verify that all fields have the expected format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>ts_hour_user_id_brand_cross</th>\n",
       "      <th>ts_weekday_user_id_brand_cross</th>\n",
       "      <th>cat_0_user_id_brand_cross</th>\n",
       "      <th>cat_1_user_id_brand_cross</th>\n",
       "      <th>cat_2_user_id_brand_cross</th>\n",
       "      <th>product_id_user_id_cross</th>\n",
       "      <th>brand_user_id_cross</th>\n",
       "      <th>ts_hour_user_id_cross</th>\n",
       "      <th>ts_minute_user_id_cross</th>\n",
       "      <th>price</th>\n",
       "      <th>ts_weekday</th>\n",
       "      <th>ts_day</th>\n",
       "      <th>ts_month</th>\n",
       "      <th>TE_brand_target_TE</th>\n",
       "      <th>TE_user_id_target_TE</th>\n",
       "      <th>TE_product_id_target_TE</th>\n",
       "      <th>TE_cat_2_target_TE</th>\n",
       "      <th>TE_ts_weekday_ts_day_target_TE</th>\n",
       "      <th>target</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>116291</td>\n",
       "      <td>2417380</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>58886</td>\n",
       "      <td>758745</td>\n",
       "      <td>103667</td>\n",
       "      <td>2997045</td>\n",
       "      <td>-0.543510</td>\n",
       "      <td>0.486228</td>\n",
       "      <td>-1.147871</td>\n",
       "      <td>1.314784</td>\n",
       "      <td>0.267353</td>\n",
       "      <td>0.278772</td>\n",
       "      <td>0.250943</td>\n",
       "      <td>-1.287949</td>\n",
       "      <td>0.377067</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1782599</td>\n",
       "      <td>1011068</td>\n",
       "      <td>1795033</td>\n",
       "      <td>2431459</td>\n",
       "      <td>1149553</td>\n",
       "      <td>1884830</td>\n",
       "      <td>1657164</td>\n",
       "      <td>1611758</td>\n",
       "      <td>4326754</td>\n",
       "      <td>0.083268</td>\n",
       "      <td>-0.996241</td>\n",
       "      <td>-1.032595</td>\n",
       "      <td>-0.864342</td>\n",
       "      <td>0.503051</td>\n",
       "      <td>0.460187</td>\n",
       "      <td>0.471215</td>\n",
       "      <td>0.459709</td>\n",
       "      <td>0.416642</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3015750</td>\n",
       "      <td>2315995</td>\n",
       "      <td>1144832</td>\n",
       "      <td>1788893</td>\n",
       "      <td>534351</td>\n",
       "      <td>1403497</td>\n",
       "      <td>130480</td>\n",
       "      <td>2726190</td>\n",
       "      <td>5845779</td>\n",
       "      <td>1.604248</td>\n",
       "      <td>0.486228</td>\n",
       "      <td>-1.032595</td>\n",
       "      <td>-0.666240</td>\n",
       "      <td>0.446776</td>\n",
       "      <td>0.390281</td>\n",
       "      <td>0.491561</td>\n",
       "      <td>0.459709</td>\n",
       "      <td>0.318418</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3532651</td>\n",
       "      <td>3400812</td>\n",
       "      <td>1062326</td>\n",
       "      <td>1707265</td>\n",
       "      <td>457204</td>\n",
       "      <td>742430</td>\n",
       "      <td>1890790</td>\n",
       "      <td>3193575</td>\n",
       "      <td>696477</td>\n",
       "      <td>-0.267539</td>\n",
       "      <td>1.474540</td>\n",
       "      <td>1.157649</td>\n",
       "      <td>-0.864342</td>\n",
       "      <td>0.466250</td>\n",
       "      <td>0.670846</td>\n",
       "      <td>0.360001</td>\n",
       "      <td>0.459752</td>\n",
       "      <td>0.419729</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1756714</td>\n",
       "      <td>384031</td>\n",
       "      <td>1667983</td>\n",
       "      <td>2305903</td>\n",
       "      <td>1028938</td>\n",
       "      <td>1704358</td>\n",
       "      <td>2163432</td>\n",
       "      <td>1588437</td>\n",
       "      <td>3423891</td>\n",
       "      <td>-0.309536</td>\n",
       "      <td>-1.490397</td>\n",
       "      <td>-0.686767</td>\n",
       "      <td>-0.666240</td>\n",
       "      <td>0.466250</td>\n",
       "      <td>0.392225</td>\n",
       "      <td>0.530786</td>\n",
       "      <td>0.459752</td>\n",
       "      <td>0.462514</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   ts_hour_user_id_brand_cross  ts_weekday_user_id_brand_cross  \\\n",
       "0                       116291                         2417380   \n",
       "1                      1782599                         1011068   \n",
       "2                      3015750                         2315995   \n",
       "3                      3532651                         3400812   \n",
       "4                      1756714                          384031   \n",
       "\n",
       "   cat_0_user_id_brand_cross  cat_1_user_id_brand_cross  \\\n",
       "0                          0                          0   \n",
       "1                    1795033                    2431459   \n",
       "2                    1144832                    1788893   \n",
       "3                    1062326                    1707265   \n",
       "4                    1667983                    2305903   \n",
       "\n",
       "   cat_2_user_id_brand_cross  product_id_user_id_cross  brand_user_id_cross  \\\n",
       "0                          0                     58886               758745   \n",
       "1                    1149553                   1884830              1657164   \n",
       "2                     534351                   1403497               130480   \n",
       "3                     457204                    742430              1890790   \n",
       "4                    1028938                   1704358              2163432   \n",
       "\n",
       "   ts_hour_user_id_cross  ts_minute_user_id_cross     price  ts_weekday  \\\n",
       "0                 103667                  2997045 -0.543510    0.486228   \n",
       "1                1611758                  4326754  0.083268   -0.996241   \n",
       "2                2726190                  5845779  1.604248    0.486228   \n",
       "3                3193575                   696477 -0.267539    1.474540   \n",
       "4                1588437                  3423891 -0.309536   -1.490397   \n",
       "\n",
       "     ts_day  ts_month  TE_brand_target_TE  TE_user_id_target_TE  \\\n",
       "0 -1.147871  1.314784            0.267353              0.278772   \n",
       "1 -1.032595 -0.864342            0.503051              0.460187   \n",
       "2 -1.032595 -0.666240            0.446776              0.390281   \n",
       "3  1.157649 -0.864342            0.466250              0.670846   \n",
       "4 -0.686767 -0.666240            0.466250              0.392225   \n",
       "\n",
       "   TE_product_id_target_TE  TE_cat_2_target_TE  \\\n",
       "0                 0.250943           -1.287949   \n",
       "1                 0.471215            0.459709   \n",
       "2                 0.491561            0.459709   \n",
       "3                 0.360001            0.459752   \n",
       "4                 0.530786            0.459752   \n",
       "\n",
       "   TE_ts_weekday_ts_day_target_TE  target  \n",
       "0                        0.377067     0.0  \n",
       "1                        0.416642     1.0  \n",
       "2                        0.318418     0.0  \n",
       "3                        0.419729     1.0  \n",
       "4                        0.462514     0.0  "
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nvtdata = pd.read_parquet(output_train_dir)\n",
    "nvtdata.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>ts_hour_user_id_brand_cross</th>\n",
       "      <th>ts_weekday_user_id_brand_cross</th>\n",
       "      <th>cat_0_user_id_brand_cross</th>\n",
       "      <th>cat_1_user_id_brand_cross</th>\n",
       "      <th>cat_2_user_id_brand_cross</th>\n",
       "      <th>product_id_user_id_cross</th>\n",
       "      <th>brand_user_id_cross</th>\n",
       "      <th>ts_hour_user_id_cross</th>\n",
       "      <th>ts_minute_user_id_cross</th>\n",
       "      <th>price</th>\n",
       "      <th>ts_weekday</th>\n",
       "      <th>ts_day</th>\n",
       "      <th>ts_month</th>\n",
       "      <th>TE_brand_target_TE</th>\n",
       "      <th>TE_user_id_target_TE</th>\n",
       "      <th>TE_product_id_target_TE</th>\n",
       "      <th>TE_cat_2_target_TE</th>\n",
       "      <th>TE_ts_weekday_ts_day_target_TE</th>\n",
       "      <th>target</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.173451</td>\n",
       "      <td>-0.502085</td>\n",
       "      <td>-1.378423</td>\n",
       "      <td>-0.468138</td>\n",
       "      <td>0.466540</td>\n",
       "      <td>0.390281</td>\n",
       "      <td>0.473820</td>\n",
       "      <td>0.459709</td>\n",
       "      <td>0.423359</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>2181500</td>\n",
       "      <td>2619444</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>-0.760359</td>\n",
       "      <td>1.474540</td>\n",
       "      <td>-1.724251</td>\n",
       "      <td>-0.468138</td>\n",
       "      <td>0.226113</td>\n",
       "      <td>0.390281</td>\n",
       "      <td>0.293626</td>\n",
       "      <td>-1.287531</td>\n",
       "      <td>0.389188</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.941599</td>\n",
       "      <td>-0.502085</td>\n",
       "      <td>0.235441</td>\n",
       "      <td>-0.468138</td>\n",
       "      <td>0.446776</td>\n",
       "      <td>0.300216</td>\n",
       "      <td>0.431577</td>\n",
       "      <td>-1.281800</td>\n",
       "      <td>0.388286</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>-0.829099</td>\n",
       "      <td>-0.996241</td>\n",
       "      <td>-0.686767</td>\n",
       "      <td>-0.468138</td>\n",
       "      <td>0.338112</td>\n",
       "      <td>0.390281</td>\n",
       "      <td>0.362127</td>\n",
       "      <td>0.305544</td>\n",
       "      <td>0.374937</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>-0.795709</td>\n",
       "      <td>-0.996241</td>\n",
       "      <td>-0.686767</td>\n",
       "      <td>-0.468138</td>\n",
       "      <td>0.364697</td>\n",
       "      <td>0.390281</td>\n",
       "      <td>0.275820</td>\n",
       "      <td>-1.287949</td>\n",
       "      <td>0.377119</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   ts_hour_user_id_brand_cross  ts_weekday_user_id_brand_cross  \\\n",
       "0                            0                               0   \n",
       "1                            0                               0   \n",
       "2                            0                               0   \n",
       "3                            0                               0   \n",
       "4                            0                               0   \n",
       "\n",
       "   cat_0_user_id_brand_cross  cat_1_user_id_brand_cross  \\\n",
       "0                          0                          0   \n",
       "1                          0                          0   \n",
       "2                          0                          0   \n",
       "3                          0                          0   \n",
       "4                          0                          0   \n",
       "\n",
       "   cat_2_user_id_brand_cross  product_id_user_id_cross  brand_user_id_cross  \\\n",
       "0                          0                         0                    0   \n",
       "1                          0                   2181500              2619444   \n",
       "2                          0                         0                    0   \n",
       "3                          0                         0                    0   \n",
       "4                          0                         0                    0   \n",
       "\n",
       "   ts_hour_user_id_cross  ts_minute_user_id_cross     price  ts_weekday  \\\n",
       "0                      0                        0  0.173451   -0.502085   \n",
       "1                      0                        0 -0.760359    1.474540   \n",
       "2                      0                        0  0.941599   -0.502085   \n",
       "3                      0                        0 -0.829099   -0.996241   \n",
       "4                      0                        0 -0.795709   -0.996241   \n",
       "\n",
       "     ts_day  ts_month  TE_brand_target_TE  TE_user_id_target_TE  \\\n",
       "0 -1.378423 -0.468138            0.466540              0.390281   \n",
       "1 -1.724251 -0.468138            0.226113              0.390281   \n",
       "2  0.235441 -0.468138            0.446776              0.300216   \n",
       "3 -0.686767 -0.468138            0.338112              0.390281   \n",
       "4 -0.686767 -0.468138            0.364697              0.390281   \n",
       "\n",
       "   TE_product_id_target_TE  TE_cat_2_target_TE  \\\n",
       "0                 0.473820            0.459709   \n",
       "1                 0.293626           -1.287531   \n",
       "2                 0.431577           -1.281800   \n",
       "3                 0.362127            0.305544   \n",
       "4                 0.275820           -1.287949   \n",
       "\n",
       "   TE_ts_weekday_ts_day_target_TE  target  \n",
       "0                        0.423359     1.0  \n",
       "1                        0.389188     0.0  \n",
       "2                        0.388286     0.0  \n",
       "3                        0.374937     1.0  \n",
       "4                        0.377119     0.0  "
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nvtdata_valid = pd.read_parquet(output_valid_dir)\n",
    "nvtdata_valid.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2359020"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sum(nvtdata_valid['ts_hour_user_id_brand_cross']==0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2461719"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(nvtdata_valid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Getting the embedding size\n",
    "\n",
    "Next, we need to get the embedding size for the categorical variables. This is an important input for defining the embedding table size to be used by HugeCTR."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'brand_user_id_cross': (3009092, 512),\n",
       " 'cat_0_user_id_brand_cross': (2877223, 512),\n",
       " 'cat_1_user_id_brand_cross': (2890639, 512),\n",
       " 'cat_2_user_id_brand_cross': (2159304, 512),\n",
       " 'product_id_user_id_cross': (4398425, 512),\n",
       " 'ts_hour_user_id_brand_cross': (4427037, 512),\n",
       " 'ts_hour_user_id_cross': (3999369, 512),\n",
       " 'ts_minute_user_id_cross': (5931061, 512),\n",
       " 'ts_weekday_user_id_brand_cross': (3961156, 512)}"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embeddings = ops.get_embedding_sizes(proc)\n",
    "embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[4427037, 3961156, 2877223, 2890639, 2159304, 4398425, 3009092, 3999369, 5931061]\n"
     ]
    }
   ],
   "source": [
    "print([embeddings[x][0] for x in cat_feats.columns])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['ts_hour_user_id_brand_cross',\n",
       " 'ts_weekday_user_id_brand_cross',\n",
       " 'cat_0_user_id_brand_cross',\n",
       " 'cat_1_user_id_brand_cross',\n",
       " 'cat_2_user_id_brand_cross',\n",
       " 'product_id_user_id_cross',\n",
       " 'brand_user_id_cross',\n",
       " 'ts_hour_user_id_cross',\n",
       " 'ts_minute_user_id_cross']"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cat_feats.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'[4427037, 3961156, 2877223, 2890639, 2159304, 4398425, 3009092, 3999369, 5931061]'"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embedding_size_str = \"{}\".format([embeddings[x][0] for x in cat_feats.columns])\n",
    "embedding_size_str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num_con_feates = len(cont_feats.columns)\n",
    "num_con_feates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['price',\n",
       " 'ts_weekday',\n",
       " 'ts_day',\n",
       " 'ts_month',\n",
       " 'TE_brand_target_TE',\n",
       " 'TE_user_id_target_TE',\n",
       " 'TE_product_id_target_TE',\n",
       " 'TE_cat_2_target_TE',\n",
       " 'TE_ts_weekday_ts_day_target_TE']"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cont_feats.columns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preparing the training configuration file for HugeCTR\n",
    "\n",
    "HugeCTR model is defined by a JSON file. The below cell defines a JSON object which contain the configuration for a DLRM network, which can then be dumped to a JSON config file. \n",
    "\n",
    "Several parameters that need to be edited to match this dataset are:\n",
    "\n",
    "- `slot_size_array`: cadinalities for the categorical variables\n",
    "- `dense_dim`: number of dense features\n",
    "- `slot_num`: number of categorical variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "dlrm_config = {\n",
    "    \"solver\": {\n",
    "      \"lr_policy\": \"fixed\",\n",
    "      \"display\": 1000,\n",
    "      \"max_iter\":1000000,\n",
    "      \"gpu\": [0,1,2,3,4,5,7],\n",
    "      \"batchsize\": 16384,\n",
    "      \"snapshot\": 10000000,\n",
    "      \"snapshot_prefix\": \"./\",\n",
    "      \"eval_interval\": 3200,\n",
    "      \"eval_batches\": 2720,\n",
    "      \"eval_metrics\": [\"AUC:1.0\"],\n",
    "      \"input_key_type\": \"I64\"\n",
    "    },\n",
    "\n",
    "    \"optimizer\": {\n",
    "       \"type\": \"SGD\",\n",
    "       \"global_update\": False,\n",
    "       \"sgd_hparam\": {\n",
    "         \"learning_rate\": 0.1,\n",
    "         \"atomic_update\": True,\n",
    "         \"warmup_steps\": 8000,\n",
    "         \"decay_start\": 48000,\n",
    "         \"decay_steps\": 24000\n",
    "       }\n",
    "     },\n",
    "\n",
    "    \"layers\": [\n",
    "        {\n",
    "        \"name\": \"data\",\n",
    "        \"type\": \"Data\",\n",
    "        \"format\": \"Parquet\",\n",
    "        \"slot_size_array\":[embeddings[x][0] for x in cat_feats.columns],\n",
    "        \"source\": \"./nvtabular_temp/output/train/_file_list.txt\",\n",
    "        \"eval_source\": \"./nvtabular_temp/output/valid/_file_list.txt\",\n",
    "        \"check\": \"None\",\n",
    "        \"label\": {\n",
    "                \"top\": \"label\",\n",
    "                \"label_dim\": 1\n",
    "        },\n",
    "        \"dense\": {\n",
    "                \"top\": \"dense\",\n",
    "                \"dense_dim\": len(cont_feats.columns)\n",
    "        },\n",
    "        \"sparse\": [\n",
    "                {\n",
    "            \"top\": \"data1\",\n",
    "            \"type\": \"DistributedSlot\",\n",
    "            \"max_feature_num_per_sample\": len(cat_feats.columns),\n",
    "            \"slot_num\": len(cat_feats.columns)\n",
    "                }\n",
    "        ]\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"sparse_embedding1\",\n",
    "        \"type\": \"DistributedSlotSparseEmbeddingHash\",\n",
    "        \"bottom\": \"data1\",\n",
    "        \"top\": \"sparse_embedding1\",\n",
    "        \"sparse_embedding_hparam\": {\n",
    "          \"embedding_vec_size\": 64,\n",
    "          \"max_vocabulary_size_per_gpu\": 5000000,\n",
    "          \"combiner\": 0\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc1\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"dense\",\n",
    "        \"top\": \"fc1\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 512\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu1\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc1\",\n",
    "        \"top\": \"relu1\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc2\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"relu1\",\n",
    "        \"top\": \"fc2\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 256\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu2\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc2\",\n",
    "        \"top\": \"relu2\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc3\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"relu2\",\n",
    "        \"top\": \"fc3\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 128\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu3\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc3\",\n",
    "        \"top\": \"relu3\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"interaction1\",\n",
    "        \"type\": \"Interaction\",\n",
    "        \"bottom\": [\"relu3\", \"sparse_embedding1\"],\n",
    "        \"top\": \"interaction1\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc4\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"interaction1\",\n",
    "        \"top\": \"fc4\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 1024\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu4\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc4\",\n",
    "        \"top\": \"relu4\"\n",
    "      },\n",
    "\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc5\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"relu4\",\n",
    "        \"top\": \"fc5\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 1024\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu5\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc5\",\n",
    "        \"top\": \"relu5\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc6\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"relu5\",\n",
    "        \"top\": \"fc6\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 512\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu6\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc6\",\n",
    "        \"top\": \"relu6\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc7\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"relu6\",\n",
    "        \"top\": \"fc7\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 256\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu7\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc7\",\n",
    "        \"top\": \"relu7\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc8\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"relu7\",\n",
    "        \"top\": \"fc8\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 1\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"loss\",\n",
    "        \"type\": \"BinaryCrossEntropyLoss\",\n",
    "        \"bottom\": [\"fc8\",\"label\"],\n",
    "        \"top\": \"loss\"\n",
    "      }\n",
    "    ]\n",
    "  }\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "with open('dlrm_config.json', 'w') as fp:\n",
    "    json.dump(dlrm_config, fp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, we can also manually write a config file, using the information obtained from the NVTabular workflow object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting dlrm_config.json\n"
     ]
    }
   ],
   "source": [
    "%%writefile dlrm_config.json\n",
    "{\n",
    "    \"solver\": {\n",
    "      \"lr_policy\": \"fixed\",\n",
    "      \"display\": 1000,\n",
    "      \"max_iter\":1000000,\n",
    "      \"gpu\": [0,1,2,3,4,5,7],\n",
    "      \"batchsize\": 16384,\n",
    "      \"snapshot\": 10000000,\n",
    "      \"snapshot_prefix\": \"./\",\n",
    "      \"eval_interval\": 3200,\n",
    "      \"eval_batches\": 2720,\n",
    "      \"eval_metrics\": [\"AUC:1.0\"],\n",
    "      \"input_key_type\": \"I64\"\n",
    "    },\n",
    "\n",
    "    \"optimizer\": {\n",
    "       \"type\": \"SGD\",\n",
    "       \"global_update\": false,\n",
    "       \"sgd_hparam\": {\n",
    "         \"learning_rate\": 0.1,\n",
    "         \"atomic_update\": true,\n",
    "         \"warmup_steps\": 8000,\n",
    "         \"decay_start\": 48000,\n",
    "         \"decay_steps\": 24000\n",
    "       }\n",
    "     },\n",
    "\n",
    "    \"layers\": [\n",
    "        {\n",
    "        \"name\": \"data\",\n",
    "        \"type\": \"Data\",\n",
    "        \"format\": \"Parquet\",\n",
    "        \"slot_size_array\":[4427037, 3961156, 2877223, 2890639, 2159304, 4398425, 3009092, 3999369, 5931061],\n",
    "        \"source\": \"./nvtabular_temp/output/train/_file_list.txt\",\n",
    "        \"eval_source\": \"./nvtabular_temp/output/valid/_file_list.txt\",\n",
    "        \"check\": \"None\",\n",
    "        \"label\": {\n",
    "                \"top\": \"label\",\n",
    "                \"label_dim\": 1\n",
    "        },\n",
    "        \"dense\": {\n",
    "                \"top\": \"dense\",\n",
    "                \"dense_dim\": 9\n",
    "        },\n",
    "        \"sparse\": [\n",
    "                {\n",
    "            \"top\": \"data1\",\n",
    "            \"type\": \"DistributedSlot\",\n",
    "            \"max_feature_num_per_sample\": 9,\n",
    "            \"slot_num\": 9\n",
    "                }\n",
    "        ]\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"sparse_embedding1\",\n",
    "        \"type\": \"DistributedSlotSparseEmbeddingHash\",\n",
    "        \"bottom\": \"data1\",\n",
    "        \"top\": \"sparse_embedding1\",\n",
    "        \"sparse_embedding_hparam\": {\n",
    "          \"embedding_vec_size\": 128,\n",
    "          \"max_vocabulary_size_per_gpu\": 10000000,\n",
    "          \"combiner\": 0\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc1\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"dense\",\n",
    "        \"top\": \"fc1\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 512\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu1\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc1\",\n",
    "        \"top\": \"relu1\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc2\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"relu1\",\n",
    "        \"top\": \"fc2\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 256\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu2\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc2\",\n",
    "        \"top\": \"relu2\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc3\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"relu2\",\n",
    "        \"top\": \"fc3\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 128\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu3\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc3\",\n",
    "        \"top\": \"relu3\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"interaction1\",\n",
    "        \"type\": \"Interaction\",\n",
    "        \"bottom\": [\"relu3\", \"sparse_embedding1\"],\n",
    "        \"top\": \"interaction1\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc4\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"interaction1\",\n",
    "        \"top\": \"fc4\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 1024\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu4\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc4\",\n",
    "        \"top\": \"relu4\"\n",
    "      },\n",
    "\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc5\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"relu4\",\n",
    "        \"top\": \"fc5\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 1024\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu5\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc5\",\n",
    "        \"top\": \"relu5\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc6\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"relu5\",\n",
    "        \"top\": \"fc6\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 512\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu6\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc6\",\n",
    "        \"top\": \"relu6\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc7\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"relu6\",\n",
    "        \"top\": \"fc7\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 256\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu7\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc7\",\n",
    "        \"top\": \"relu7\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc8\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"relu7\",\n",
    "        \"top\": \"fc8\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 1\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"loss\",\n",
    "        \"type\": \"BinaryCrossEntropyLoss\",\n",
    "        \"bottom\": [\"fc8\",\"label\"],\n",
    "        \"top\": \"loss\"\n",
    "      }\n",
    "    ]\n",
    "  }\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll shutdown our Dask client from earlier to free up some memory so that we can share it with HugeCTR."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.shutdown()\n",
    "cluster.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"3\"></a>\n",
    "## 3. HugeCTR training\n",
    "\n",
    "Now we are ready to train a DLRM model with HugeCTR.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "# HugeCTR\n",
    "import sys\n",
    "sys.path.append(\"/usr/local/hugectr/lib\")\n",
    "\n",
    "import os\n",
    "from mpi4py import MPI\n",
    "os.environ['NCCL_SHM_DISABLE']='1'\n",
    "from hugectr import Session, solver_parser_helper, get_learning_rate_scheduler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "distributed.client - ERROR - Failed to reconnect to scheduler after 10.00 seconds, closing client\n",
      "_GatheringFuture exception was never retrieved\n",
      "future: <_GatheringFuture finished exception=CancelledError()>\n",
      "asyncio.exceptions.CancelledError\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[HUGECTR][INFO] iter: 0; loss: 0.6649990081787109\n",
      "[HUGECTR][INFO] iter: 100; loss: 0.6585077047348022\n",
      "[HUGECTR][INFO] iter: 200; loss: 0.6553144454956055\n",
      "[HUGECTR][INFO] iter: 300; loss: 0.6533859968185425\n",
      "[HUGECTR][INFO] iter: 400; loss: 0.6496698260307312\n",
      "[HUGECTR][INFO] iter: 500; loss: 0.6550267338752747\n",
      "[HUGECTR][INFO] iter: 600; loss: 0.6582932472229004\n",
      "[HUGECTR][INFO] iter: 700; loss: 0.6601693034172058\n",
      "[HUGECTR][INFO] iter: 800; loss: 0.6651085615158081\n",
      "[HUGECTR][INFO] iter: 900; loss: 0.660534679889679\n",
      "[HUGECTR][INFO] iter: 1000; loss: 0.6576811075210571\n",
      "[HUGECTR][INFO] iter: 1100; loss: 0.6517642140388489\n",
      "[HUGECTR][INFO] iter: 1200; loss: 0.6499871015548706\n",
      "[HUGECTR][INFO] iter: 1300; loss: 0.6492666006088257\n",
      "[HUGECTR][INFO] iter: 1400; loss: 0.6422186493873596\n",
      "[HUGECTR][INFO] iter: 1500; loss: 0.6258538961410522\n",
      "[HUGECTR][INFO] iter: 1600; loss: 0.5962243676185608\n",
      "[HUGECTR][INFO] iter: 1700; loss: 0.5636842250823975\n",
      "[HUGECTR][INFO] iter: 1800; loss: 0.5577143430709839\n",
      "[HUGECTR][INFO] iter: 1900; loss: 0.5541603565216064\n",
      "[HUGECTR][INFO] iter: 2000; loss: 0.5259717106819153\n",
      "[HUGECTR][INFO] iter: 2100; loss: 0.5111467838287354\n",
      "[HUGECTR][INFO] iter: 2200; loss: 0.5102323293685913\n",
      "[HUGECTR][INFO] iter: 2300; loss: 0.508430540561676\n",
      "[HUGECTR][INFO] iter: 2400; loss: 0.5129308700561523\n",
      "[HUGECTR][INFO] iter: 2500; loss: 0.5244858264923096\n",
      "[HUGECTR][INFO] iter: 2600; loss: 0.5238444209098816\n",
      "[HUGECTR][INFO] iter: 2700; loss: 0.521399736404419\n",
      "[HUGECTR][INFO] iter: 2800; loss: 0.53584885597229\n",
      "[HUGECTR][INFO] iter: 2900; loss: 0.5296216607093811\n",
      "[HUGECTR][INFO] iter: 3000; loss: 0.5244985818862915\n",
      "[HUGECTR][INFO] iter: 3000, [('AUC', 0.6373581886291504)]\n",
      "[HUGECTR][INFO] iter: 3100; loss: 0.5340688824653625\n",
      "[HUGECTR][INFO] iter: 3200; loss: 0.5352561473846436\n",
      "[HUGECTR][INFO] iter: 3300; loss: 0.5306345224380493\n",
      "[HUGECTR][INFO] iter: 3400; loss: 0.5279024243354797\n",
      "[HUGECTR][INFO] iter: 3500; loss: 0.5378730893135071\n",
      "[HUGECTR][INFO] iter: 3600; loss: 0.5347540974617004\n",
      "[HUGECTR][INFO] iter: 3700; loss: 0.5299892425537109\n",
      "[HUGECTR][INFO] iter: 3800; loss: 0.5212597250938416\n",
      "[HUGECTR][INFO] iter: 3900; loss: 0.4974917769432068\n",
      "[HUGECTR][INFO] iter: 4000; loss: 0.49960947036743164\n",
      "[HUGECTR][INFO] iter: 4100; loss: 0.4960770606994629\n",
      "[HUGECTR][INFO] iter: 4200; loss: 0.5081725120544434\n",
      "[HUGECTR][INFO] iter: 4300; loss: 0.5063897371292114\n",
      "[HUGECTR][INFO] iter: 4400; loss: 0.5171743631362915\n",
      "[HUGECTR][INFO] iter: 4500; loss: 0.5146521925926208\n",
      "[HUGECTR][INFO] iter: 4600; loss: 0.5158137679100037\n",
      "[HUGECTR][INFO] iter: 4700; loss: 0.5238028764724731\n",
      "[HUGECTR][INFO] iter: 4800; loss: 0.5270493626594543\n",
      "[HUGECTR][INFO] iter: 4900; loss: 0.5186280012130737\n",
      "[HUGECTR][INFO] iter: 5000; loss: 0.522905170917511\n",
      "[HUGECTR][INFO] iter: 5100; loss: 0.5267692804336548\n",
      "[HUGECTR][INFO] iter: 5200; loss: 0.5324996709823608\n",
      "[HUGECTR][INFO] iter: 5300; loss: 0.5294499397277832\n",
      "[HUGECTR][INFO] iter: 5400; loss: 0.5299464464187622\n",
      "[HUGECTR][INFO] iter: 5500; loss: 0.5394059419631958\n",
      "[HUGECTR][INFO] iter: 5600; loss: 0.5341923832893372\n",
      "[HUGECTR][INFO] iter: 5700; loss: 0.5235053896903992\n",
      "[HUGECTR][INFO] iter: 5800; loss: 0.528168797492981\n",
      "[HUGECTR][INFO] iter: 5900; loss: 0.49902039766311646\n",
      "[HUGECTR][INFO] iter: 6000; loss: 0.48950284719467163\n",
      "[HUGECTR][INFO] iter: 6000, [('AUC', 0.6533010005950928)]\n",
      "[HUGECTR][INFO] iter: 6100; loss: 0.5026806592941284\n",
      "[HUGECTR][INFO] iter: 6200; loss: 0.5064517855644226\n",
      "[HUGECTR][INFO] iter: 6300; loss: 0.5059552192687988\n",
      "[HUGECTR][INFO] iter: 6400; loss: 0.5104296207427979\n",
      "[HUGECTR][INFO] iter: 6500; loss: 0.5119836926460266\n",
      "[HUGECTR][INFO] iter: 6600; loss: 0.5268381834030151\n",
      "[HUGECTR][INFO] iter: 6700; loss: 0.5214208960533142\n",
      "[HUGECTR][INFO] iter: 6800; loss: 0.5157287120819092\n",
      "[HUGECTR][INFO] iter: 6900; loss: 0.5147594213485718\n",
      "[HUGECTR][INFO] iter: 7000; loss: 0.5171025991439819\n",
      "[HUGECTR][INFO] iter: 7100; loss: 0.5233388543128967\n",
      "[HUGECTR][INFO] iter: 7200; loss: 0.5234604477882385\n",
      "[HUGECTR][INFO] iter: 7300; loss: 0.5271459817886353\n",
      "[HUGECTR][INFO] iter: 7400; loss: 0.5290958881378174\n",
      "[HUGECTR][INFO] iter: 7500; loss: 0.5344234704971313\n",
      "[HUGECTR][INFO] iter: 7600; loss: 0.5346677303314209\n",
      "[HUGECTR][INFO] iter: 7700; loss: 0.5212712287902832\n",
      "[HUGECTR][INFO] iter: 7800; loss: 0.49437257647514343\n",
      "[HUGECTR][INFO] iter: 7900; loss: 0.5086170434951782\n",
      "[HUGECTR][INFO] iter: 8000; loss: 0.4940020740032196\n",
      "[HUGECTR][INFO] iter: 8100; loss: 0.49439147114753723\n",
      "[HUGECTR][INFO] iter: 8200; loss: 0.5066720247268677\n",
      "[HUGECTR][INFO] iter: 8300; loss: 0.4989408254623413\n",
      "[HUGECTR][INFO] iter: 8400; loss: 0.5073046088218689\n",
      "[HUGECTR][INFO] iter: 8500; loss: 0.5048984289169312\n",
      "[HUGECTR][INFO] iter: 8600; loss: 0.5138943791389465\n",
      "[HUGECTR][INFO] iter: 8700; loss: 0.5158380270004272\n",
      "[HUGECTR][INFO] iter: 8800; loss: 0.5127506256103516\n",
      "[HUGECTR][INFO] iter: 8900; loss: 0.5110788345336914\n",
      "[HUGECTR][INFO] iter: 9000; loss: 0.5161044597625732\n",
      "[HUGECTR][INFO] iter: 9000, [('AUC', 0.6372886896133423)]\n",
      "[HUGECTR][INFO] iter: 9100; loss: 0.5157985091209412\n",
      "[HUGECTR][INFO] iter: 9200; loss: 0.5275703072547913\n",
      "[HUGECTR][INFO] iter: 9300; loss: 0.5305440425872803\n",
      "[HUGECTR][INFO] iter: 9400; loss: 0.522693395614624\n",
      "[HUGECTR][INFO] iter: 9500; loss: 0.5335985422134399\n",
      "[HUGECTR][INFO] iter: 9600; loss: 0.5230779051780701\n",
      "[HUGECTR][INFO] iter: 9700; loss: 0.5202857255935669\n",
      "[HUGECTR][INFO] iter: 9800; loss: 0.49916335940361023\n",
      "[HUGECTR][INFO] iter: 9900; loss: 0.49024444818496704\n",
      "CPU times: user 40min 43s, sys: 1h 13min 22s, total: 1h 54min 5s\n",
      "Wall time: 23min 20s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Error_t.Success"
      ]
     },
     "execution_count": 71,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "#NUM_GPUS = [0,1,2,3,4,5,6,7]\n",
    "NUM_GPUS = [0,1,2,3]\n",
    "\n",
    "# Set config file\n",
    "json_file = \"dlrm_config.json\"\n",
    "# Set solver config\n",
    "solver_config = solver_parser_helper(seed = 0,\n",
    "                                     batchsize = 16384,\n",
    "                                     batchsize_eval = 16384,\n",
    "                                     model_file = \"\",\n",
    "                                     embedding_files = [],\n",
    "                                     vvgpu = [NUM_GPUS],\n",
    "                                     use_mixed_precision = False,\n",
    "                                     scaler = 1.0,\n",
    "                                     i64_input_key = True,\n",
    "                                     use_algorithm_search = True,\n",
    "                                     use_cuda_graph = True,\n",
    "                                     repeat_dataset = True\n",
    "                                    )\n",
    "# Set learning rate\n",
    "lr_sch = get_learning_rate_scheduler(json_file)\n",
    "# Train model\n",
    "sess = Session(solver_config, json_file)\n",
    "sess.start_data_reading()\n",
    "for i in range(10000):\n",
    "    lr = lr_sch.get_next()\n",
    "    sess.set_learning_rate(lr)\n",
    "    sess.train()\n",
    "    if (i%100 == 0):\n",
    "        loss = sess.get_current_loss()\n",
    "        print(\"[HUGECTR][INFO] iter: {}; loss: {}\".format(i, loss))\n",
    "    if (i%3000 == 0 and i != 0):\n",
    "        sess.check_overflow()\n",
    "        sess.copy_weights_for_evaluation()\n",
    "        metrics = sess.evaluation()\n",
    "        print(\"[HUGECTR][INFO] iter: {}, {}\".format(i, metrics))\n",
    "sess.download_params_to_files(\"./\", i+1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"4\"></a>\n",
    "## 4. HugeCTR inference\n",
    "\n",
    "In this section, we will read the test dataset, and compute the AUC value. \n",
    "\n",
    "First, we need to prepare a JSON file for inference. Several modification to the training config is required.\n",
    "\n",
    "- We need to omit the `optimizer` and `solver` clauses, while adding and `inference` clause\n",
    "- We need to change the output layer to `Sigmoid` type. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting dlrm_inference.json\n"
     ]
    }
   ],
   "source": [
    "%%writefile dlrm_inference.json\n",
    "{\n",
    "  \"inference\": {\n",
    "    \"max_batchsize\": 4096,\n",
    "    \"dense_model_file\": \"./_dense_10000.model\",\n",
    "    \"sparse_model_file\": \"./0_sparse_10000.model\",\n",
    "    \"input_key_type\": \"I64\"\n",
    "  },\n",
    "    \"layers\": [\n",
    "        {\n",
    "        \"name\": \"data\",\n",
    "        \"type\": \"Data\",\n",
    "        \"label\": {\n",
    "                \"top\": \"label\",\n",
    "                \"label_dim\": 1\n",
    "        },\n",
    "        \"dense\": {\n",
    "                \"top\": \"dense\",\n",
    "                \"dense_dim\": 9\n",
    "        },\n",
    "        \"sparse\": [\n",
    "                {\n",
    "            \"top\": \"data1\",\n",
    "            \"type\": \"DistributedSlot\",\n",
    "            \"max_feature_num_per_sample\": 9,\n",
    "            \"slot_num\": 9\n",
    "                }\n",
    "        ]\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"sparse_embedding1\",\n",
    "        \"type\": \"DistributedSlotSparseEmbeddingHash\",\n",
    "        \"bottom\": \"data1\",\n",
    "        \"top\": \"sparse_embedding1\",\n",
    "        \"sparse_embedding_hparam\": {\n",
    "          \"embedding_vec_size\": 128,\n",
    "          \"max_vocabulary_size_per_gpu\": 5000000,\n",
    "          \"combiner\": 0\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc1\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"dense\",\n",
    "        \"top\": \"fc1\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 512\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu1\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc1\",\n",
    "        \"top\": \"relu1\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc2\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"relu1\",\n",
    "        \"top\": \"fc2\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 256\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu2\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc2\",\n",
    "        \"top\": \"relu2\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc3\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"relu2\",\n",
    "        \"top\": \"fc3\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 128\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu3\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc3\",\n",
    "        \"top\": \"relu3\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"interaction1\",\n",
    "        \"type\": \"Interaction\",\n",
    "        \"bottom\": [\"relu3\", \"sparse_embedding1\"],\n",
    "        \"top\": \"interaction1\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc4\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"interaction1\",\n",
    "        \"top\": \"fc4\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 1024\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu4\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc4\",\n",
    "        \"top\": \"relu4\"\n",
    "      },\n",
    "\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc5\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"relu4\",\n",
    "        \"top\": \"fc5\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 1024\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu5\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc5\",\n",
    "        \"top\": \"relu5\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc6\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"relu5\",\n",
    "        \"top\": \"fc6\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 512\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu6\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc6\",\n",
    "        \"top\": \"relu6\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc7\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"relu6\",\n",
    "        \"top\": \"fc7\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 256\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"relu7\",\n",
    "        \"type\": \"ReLU\",\n",
    "        \"bottom\": \"fc7\",\n",
    "        \"top\": \"relu7\"\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"fc8\",\n",
    "        \"type\": \"InnerProduct\",\n",
    "        \"bottom\": \"relu7\",\n",
    "        \"top\": \"fc8\",\n",
    "         \"fc_param\": {\n",
    "          \"num_output\": 1\n",
    "        }\n",
    "      },\n",
    "\n",
    "      {\n",
    "        \"name\": \"sigmoid\",\n",
    "        \"type\": \"Sigmoid\",\n",
    "        \"bottom\": [\"fc8\"],\n",
    "        \"top\": \"sigmoid\"\n",
    "      }\n",
    "    ]\n",
    "  }\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare the inference session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "from mpi4py import MPI\n",
    "from hugectr.inference import CreateParameterServer, CreateEmbeddingCache, InferenceSession\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create parameter server, embedding cache and inference session\n",
    "parameter_server = CreateParameterServer([\"dlrm_inference.json\"], [\"DLRM\"], True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_cache = CreateEmbeddingCache(parameter_server, 0, True, 0.2, \"dlrm_inference.json\", \"DLRM\", True)\n",
    "inference_session = InferenceSession(\"dlrm_inference.json\", 0, embedding_cache)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reading and prepare the data\n",
    "\n",
    "We first read the NVTabular processed data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>ts_hour_user_id_brand_cross</th>\n",
       "      <th>ts_weekday_user_id_brand_cross</th>\n",
       "      <th>cat_0_user_id_brand_cross</th>\n",
       "      <th>cat_1_user_id_brand_cross</th>\n",
       "      <th>cat_2_user_id_brand_cross</th>\n",
       "      <th>product_id_user_id_cross</th>\n",
       "      <th>brand_user_id_cross</th>\n",
       "      <th>ts_hour_user_id_cross</th>\n",
       "      <th>ts_minute_user_id_cross</th>\n",
       "      <th>price</th>\n",
       "      <th>ts_weekday</th>\n",
       "      <th>ts_day</th>\n",
       "      <th>ts_month</th>\n",
       "      <th>TE_brand_target_TE</th>\n",
       "      <th>TE_user_id_target_TE</th>\n",
       "      <th>TE_product_id_target_TE</th>\n",
       "      <th>TE_cat_2_target_TE</th>\n",
       "      <th>TE_ts_weekday_ts_day_target_TE</th>\n",
       "      <th>target</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>-0.074429</td>\n",
       "      <td>-0.007929</td>\n",
       "      <td>1.618752</td>\n",
       "      <td>-0.270035</td>\n",
       "      <td>0.334199</td>\n",
       "      <td>0.408567</td>\n",
       "      <td>0.300926</td>\n",
       "      <td>0.341582</td>\n",
       "      <td>0.239161</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>-0.326691</td>\n",
       "      <td>-1.490397</td>\n",
       "      <td>1.272925</td>\n",
       "      <td>-0.270035</td>\n",
       "      <td>0.466540</td>\n",
       "      <td>0.390281</td>\n",
       "      <td>0.531653</td>\n",
       "      <td>0.459709</td>\n",
       "      <td>0.444884</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.318253</td>\n",
       "      <td>0.980384</td>\n",
       "      <td>-1.378423</td>\n",
       "      <td>-0.270035</td>\n",
       "      <td>0.441240</td>\n",
       "      <td>0.390281</td>\n",
       "      <td>0.516564</td>\n",
       "      <td>0.423564</td>\n",
       "      <td>0.418208</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>3288133</td>\n",
       "      <td>0</td>\n",
       "      <td>1.155723</td>\n",
       "      <td>-0.502085</td>\n",
       "      <td>-1.724251</td>\n",
       "      <td>-0.270035</td>\n",
       "      <td>0.446606</td>\n",
       "      <td>0.432225</td>\n",
       "      <td>0.474201</td>\n",
       "      <td>0.459691</td>\n",
       "      <td>0.063315</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0.697609</td>\n",
       "      <td>-0.007929</td>\n",
       "      <td>-1.608975</td>\n",
       "      <td>-0.270035</td>\n",
       "      <td>0.282870</td>\n",
       "      <td>0.390281</td>\n",
       "      <td>0.286815</td>\n",
       "      <td>0.372063</td>\n",
       "      <td>0.000137</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   ts_hour_user_id_brand_cross  ts_weekday_user_id_brand_cross  \\\n",
       "0                            0                               0   \n",
       "1                            0                               0   \n",
       "2                            0                               0   \n",
       "3                            0                               0   \n",
       "4                            0                               0   \n",
       "\n",
       "   cat_0_user_id_brand_cross  cat_1_user_id_brand_cross  \\\n",
       "0                          0                          0   \n",
       "1                          0                          0   \n",
       "2                          0                          0   \n",
       "3                          0                          0   \n",
       "4                          0                          0   \n",
       "\n",
       "   cat_2_user_id_brand_cross  product_id_user_id_cross  brand_user_id_cross  \\\n",
       "0                          0                         0                    0   \n",
       "1                          0                         0                    0   \n",
       "2                          0                         0                    0   \n",
       "3                          0                         0                    0   \n",
       "4                          0                         0                    0   \n",
       "\n",
       "   ts_hour_user_id_cross  ts_minute_user_id_cross     price  ts_weekday  \\\n",
       "0                      0                        0 -0.074429   -0.007929   \n",
       "1                      0                        0 -0.326691   -1.490397   \n",
       "2                      0                        0  0.318253    0.980384   \n",
       "3                3288133                        0  1.155723   -0.502085   \n",
       "4                      0                        0  0.697609   -0.007929   \n",
       "\n",
       "     ts_day  ts_month  TE_brand_target_TE  TE_user_id_target_TE  \\\n",
       "0  1.618752 -0.270035            0.334199              0.408567   \n",
       "1  1.272925 -0.270035            0.466540              0.390281   \n",
       "2 -1.378423 -0.270035            0.441240              0.390281   \n",
       "3 -1.724251 -0.270035            0.446606              0.432225   \n",
       "4 -1.608975 -0.270035            0.282870              0.390281   \n",
       "\n",
       "   TE_product_id_target_TE  TE_cat_2_target_TE  \\\n",
       "0                 0.300926            0.341582   \n",
       "1                 0.531653            0.459709   \n",
       "2                 0.516564            0.423564   \n",
       "3                 0.474201            0.459691   \n",
       "4                 0.286815            0.372063   \n",
       "\n",
       "   TE_ts_weekday_ts_day_target_TE  target  \n",
       "0                        0.239161     1.0  \n",
       "1                        0.444884     1.0  \n",
       "2                        0.418208     0.0  \n",
       "3                        0.063315     0.0  \n",
       "4                        0.000137     0.0  "
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "\n",
    "nvtdata_test = pd.read_parquet('./nvtabular_temp/output/test')\n",
    "nvtdata_test.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "con_feats = ['price',\n",
    " 'ts_weekday',\n",
    " 'ts_day',\n",
    " 'ts_month',\n",
    " 'TE_brand_target_TE',\n",
    " 'TE_user_id_target_TE',\n",
    " 'TE_product_id_target_TE',\n",
    " 'TE_cat_2_target_TE',\n",
    " 'TE_ts_weekday_ts_day_target_TE']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [],
   "source": [
    "cat_feats = ['ts_hour_user_id_brand_cross',\n",
    " 'ts_weekday_user_id_brand_cross',\n",
    " 'cat_0_user_id_brand_cross',\n",
    " 'cat_1_user_id_brand_cross',\n",
    " 'cat_2_user_id_brand_cross',\n",
    " 'product_id_user_id_cross',\n",
    " 'brand_user_id_cross',\n",
    " 'ts_hour_user_id_cross',\n",
    " 'ts_minute_user_id_cross']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "emb_size = [4427037, 3961156, 2877223, 2890639, 2159304, 4398425, 3009092, 3999369, 5931061]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Converting data to CSR format\n",
    "\n",
    "HugeCTR expects data in CSR format for inference. One important thing to note is that NVTabular requires categorical variables to occupy different integer ranges. For example, if there are 10 users and 10 items, then the users should be encoded in the 0-9 range, while items should be in the 10-19 range. NVTabular encodes both users and items in the 0-9 ranges.\n",
    "\n",
    "For this reason, we need to shift the keys of the categorical variable produced by NVTabular to comply with HugeCTR."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "shift = np.insert(np.cumsum(emb_size), 0, 0)[:-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [],
   "source": [
    "cat_data = nvtdata_test[cat_feats].values + shift"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [],
   "source": [
    "dense_data = nvtdata_test[con_feats].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [],
   "source": [
    "def infer_batch(inference_session, dense_data_batch, cat_data_batch):\n",
    "    dense_features = list(dense_data_batch.flatten())\n",
    "    embedding_columns = list(cat_data_batch.flatten())\n",
    "\n",
    "    row_ptrs= list(range(0,len(embedding_columns)+1))\n",
    "    \n",
    "    output = inference_session.predict(dense_features, embedding_columns, row_ptrs, True)\n",
    "    return output\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we are ready to carry out inference on the test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 4096\n",
    "num_batches = (len(dense_data) // batch_size) + 1\n",
    "batch_idx = np.array_split(np.arange(len(dense_data)), num_batches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://pypi.org/simple, https://pypi.ngc.nvidia.com\n",
      "Requirement already satisfied: tqdm in /opt/conda/envs/rapids/lib/python3.8/site-packages (4.58.0)\n"
     ]
    }
   ],
   "source": [
    "!pip install tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 677/677 [00:07<00:00, 87.50it/s]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "labels = []\n",
    "for batch_id in tqdm(batch_idx):\n",
    "    dense_data_batch = dense_data[batch_id]\n",
    "    cat_data_batch = cat_data[batch_id]\n",
    "    results = infer_batch(inference_session, dense_data_batch, cat_data_batch)\n",
    "    labels.extend(results)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2772486"
      ]
     },
     "execution_count": 87,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Computing the test AUC value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "ground_truth = nvtdata_test['target'].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.6434478680485719"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from sklearn.metrics import roc_auc_score\n",
    "\n",
    "roc_auc_score(ground_truth, labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Conclusion\n",
    "\n",
    "In this notebook, we have walked you through the process of preprocessing the data, train a DLRM model with HugeCTR, then carrying out inference with the HugeCTR Python interface. Try this workflow on your data and let us know your feedback.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
